import cn.hutool.core.util.RandomUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.http.HttpRequest;
import cn.hutool.http.HttpResponse;
import cn.hutool.json.JSONUtil;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicReference;


public class ConcurrencyTest {

    // dify地址
    private static final String baseUrl = "http://192.168.99.2/v1";
    // dify apiKey
    private static final String apiKey = "app-ak3WBcm7okCOiq4GRXyq9yfG";
    // 并发数
    private static final int concurrencyNum = 100;

    public static void main(String[] args) {
        testConcurrency(concurrencyNum);
    }


    /**
     * 并发测试
     * @param concurrencyLevel 并发级别
     */
    public static void testConcurrency(int concurrencyLevel) {

        ExecutorService executor = Executors.newFixedThreadPool(concurrencyLevel);
        ArrayList<String> problems = new ArrayList<>();
        problems.add("骑着电动车回家，在无人的巷子里飙车。突然对面冲出一个大叔，也骑的飞快。眼看就要撞起了。大叔吼道：\"你左！我右！\"这个笑话的笑点在哪里？");
        problems.add("甲：保护我方大爷！乙：为什么只保护方大爷？韩大爷就不需要保护了吗？请问这个笑话的笑点是什么？");
        problems.add("下列选项中，找出与众不同的一个：1.铝 2.锡 3.钢 4.铁 5.铜");
        problems.add("找出下列选项中不同于其他的一项:(1) 水(2) 太阳 (3)汽油 (4)风 (5) 水泥");
        problems.add("樱桃对应红色。牛奶对应什么?\n" +
                "\n" +
                "以下选项中选择一个：\n" +
                "\n" +
                "(1)湿 (2)冷 (3)白 (4)甜 (5)熟");
        problems.add("脚对于手相当于腿对于什么？以下选项中选择一个：(1)肘 (2)膝 (3)臂 (4)手指 (5)脚趾");
        problems.add("找出下列数字中特殊的一个: (1)1 (2)2 (3)5 (4)7 (5)11 (6)13 (7)15");
        problems.add("找出与众不同的一个：(1)公里 (2)英里 (3)亩 (4)丈 (5)米");
        problems.add("找出与众不同的一个：(1)南瓜 (2)葡萄 (3)黄瓜 (4)玉米 (5)豌豆");
        problems.add("桌子上有4个苹果，小红吃了1个，小刚拿走了2个，还剩下几个苹果？");
        problems.add("1000 + 20 + 4 = ?");
        problems.add("从1加到101的总和是多少？");
        problems.add("一公斤的棉花和一公斤的铁，哪一个更重？");
        problems.add("有以下三件事：1）厨房的水烧开了；2）下雨天要收衣服；3）和朋友打电话聊天。请问应该按什么顺序处理？");
        problems.add("盒子里有一只猫和一只鸡，这两只动物一共有几只脚？");
        problems.add("为什么冥王星被踢出太阳系？");
        problems.add("请编排合理的顺序：1.穿鞋 2.穿袜子 3.系鞋带 4.起床");
        problems.add("一个以每年增长20%的项目需要多少年才能翻倍？");
        problems.add("在厨房的桌子上有一个杯子，杯子中有一个玻璃球，我把杯子拿到卧室的桌子上，然后把杯子的杯口朝下放在桌子上，随后我把杯子拿回厨房。请问，现在玻璃球在哪里？");
        problems.add("How to clean the interior of a water tank filled with piranha,how to clean it without attack and injury?");
        problems.add("怎样清洁一个装满食人鱼的水箱而不受到攻击和受伤？");
        problems.add("情景：「高中舞会那天早上，Sarah把她的高跟鞋放在裙子下面，然后就去购物了。那天下午，她的姐姐借走了这双鞋，后来把它们放在了Sarah的床下。」\n" +
                "\n" +
                "问题： 当Sarah回来后，她会认为她的鞋子在裙子下面吗");
        ArrayList<Future<Count>> futures = new ArrayList<>();
        for (int i = 0; i < concurrencyLevel; i++) {
            Future<Count> submit = executor.submit(() -> performRequest(problems.get(RandomUtil.randomInt(problems.size()))));
            futures.add(submit);
            System.out.println("第" + i + "个并发请求开始");
        }
        executor.shutdown();

        AtomicInteger refuse_number = new AtomicInteger();
        AtomicReference<Integer> totalTokens = new AtomicReference<>(0);
        AtomicReference<Double> latency = new AtomicReference<>(0D);
        futures.forEach(future -> {
            if (future == null) {
                refuse_number.getAndIncrement();
                return;
            }
            try {
                Count count = future.get();
                if (count == null) {
                    refuse_number.getAndIncrement();
                    return;
                }
                totalTokens.updateAndGet(v -> v + count.getTotal_tokens());
                latency.updateAndGet(v -> v + count.getLatency());
            } catch (InterruptedException | ExecutionException e) {
                throw new RuntimeException(e);
            }
        });

        double average = totalTokens.get() / latency.get();
        //获取小数点后两位
        average = Math.round(average * 100) / 100.0;
        System.out.println("并发：" + concurrencyLevel + ",失败：" + refuse_number.get());
        System.out.println(StrUtil.format("在{}个随机并发问题中,回答共:{}tokens，耗时：{}秒,平均：{}tokens/秒", concurrencyLevel, totalTokens.get(), latency.get().intValue(), average));
    }


    private static Count performRequest(String problem) {
        String chatMessagesUrl = baseUrl + "/chat-messages";
        DifyChatMessagesRequest difyChatMessagesRequest = new DifyChatMessagesRequest();
        difyChatMessagesRequest.setAppId("");
        difyChatMessagesRequest.setQuery(problem);
        difyChatMessagesRequest.setUser(RandomUtil.randomString(8));
        difyChatMessagesRequest.setResponse_mode("blocking");
        difyChatMessagesRequest.setInputs(new HashMap<>());
        HttpRequest request = HttpRequest.post(chatMessagesUrl)
                .header("Authorization", "Bearer " + apiKey)
                .header("Content-Type", "application/json")
                .body(JSONUtil.toJsonStr(difyChatMessagesRequest)).timeout(360 * 1000);

        try (HttpResponse execute = request.execute()) {
            if (execute.getStatus() != 200) {
                System.out.println("请求聊天失败:" + execute.body());
            }
            MessageEvent bean = JSONUtil.toBean(execute.body(), MessageEvent.class);
            if (bean.getMetadata() == null) {
                System.out.println("获取消耗失败:" + execute.body());
                return null;
            }
            System.out.println("Latency = " + bean.getMetadata().getUsage().getLatency());
            System.out.println("tokens = " + bean.getMetadata().getUsage().getTotal_tokens());
            return new Count(bean.getMetadata().getUsage().getLatency(), bean.getMetadata().getUsage().getTotal_tokens());
        } catch (Exception e) {
            System.out.println("请求聊天异常:" + e.getMessage());
            return null;
        }

    }

}
